
import torch
import torch.nn as nn
import math
import numpy as np
from typing import List, Dict, Tuple

from ..common import constants

PAD_WORD = "<PAD>"
UNK_WORD = "<UNK>"
BOS_WORD = "<SOS>"
EOS_WORD = "<EOS>"

class PositionalEncoding(nn.Module):
    r"""
    为输入序列中的每个位置添加独特的编码向量。
    """
  
    def __init__(self, embedding_size: int, max_sequence_length: int):
        """
        Args:
            embedding_size: 词向量的维度
            max_sequence_length: 最长序列长度
        """
        
        super(PositionalEncoding, self).__init__()
        # 生成正弦位置编码表
        sinusoid_table = PositionalEncoding.__get_sinusoid_encoding_table(embedding_size, max_sequence_length)
        self.register_buffer("sinusoid_table", sinusoid_table)

    def forward(self, x: torch.tensor) -> torch.Tensor:
        # 获取张量 x 对应长度的位置编码
        sinusoid_encoding = self.sinusoid_table[:, :x.size(1)]
        # 对张量 x 与位置编码做相加操作
        output = x + sinusoid_encoding
        
        return output
    
    @staticmethod
    def __get_sinusoid_encoding_table(embedding_size: int, max_sequence_length: int) -> torch.Tensor:
        sinusoid_list = [PositionalEncoding.__get_position_angle_vec(embedding_size, curr_position)
                         for curr_position in range(max_sequence_length)]
        sinusoid_ndarray = np.array(sinusoid_list)
        sinusoid_ndarray[:, 0::2] = np.sin(sinusoid_ndarray[:, 0::2])
        sinusoid_ndarray[:, 1::2] = np.cos(sinusoid_ndarray[:, 1::2])
        sinusoid_tensor = torch.FloatTensor(sinusoid_ndarray)
        sinusoid_tensor = sinusoid_tensor.unsqueeze(0)

        return sinusoid_tensor
    
    @staticmethod
    def __get_position_angle_vec(embedding_size: int, curr_position: int) -> List[float]:
        result = [curr_position / math.pow(10000, 2 * (curr_dim // 2) / embedding_size)
                  for curr_dim in range(embedding_size)]

        return result
 
 
class PositionWiseFeedForward(nn.Module):
    r"""
    编码器和解码器的每层都包含一个全连接前馈网络。
    """
    
    def __init__(self, embedding_size: int, hidden_size: int, dropout_prob=constants.MODEL_DROPOUT_PROB_DEFAULT):
        """
        Args:
            embedding_size: 词向量的维度
            hidden_size: 隐藏层维度
            dropout_prob: 丢弃率, 防止过拟合
        """
        
        super(PositionWiseFeedForward, self).__init__()
        self.linear_1 = nn.Linear(embedding_size, hidden_size)
        self.linear_2 = nn.Linear(hidden_size, embedding_size)
        self.layer_norm = nn.LayerNorm(embedding_size, eps=constants.LAYER_NORM_EPS_DEFAULT)
        self.dropout = nn.Dropout(dropout_prob)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        
        output = self.linear_1(x)
        output = nn.functional.relu(output)
        # 通过丢弃法防止过拟合
        output = self.dropout(output)
        
        output = self.linear_2(output)
        # 做残差连接操作
        output = output + residual
        # 做归一化操作
        output = self.layer_norm(output)
        
        return output


class ScaledDotProductAttention(nn.Module):
    r"""
    缩放点积注意力通过计算查询 Q 与键 K 之间的相似度来分配注意力权重。
    """
    
    def __init__(self, key_dim: int, dropout_prob=constants.MODEL_DROPOUT_PROB_DEFAULT):
        """
        Args:
            key_dim: key维度
            dropout_prob: 丢弃率, 防止过拟合
        """
        
        super(ScaledDotProductAttention, self).__init__()
        self.key_dim = key_dim
        self.sqrt_of_key_dim = math.sqrt(key_dim)
        self.dropout = nn.Dropout(dropout_prob)
        
    def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, 
                attention_mask: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        # 转置前 key 的 shape 属性是 (批次大小, 多头注意力的数量, 序列长度, key维度),
        # 转置后 key_t 的 shape 属性是 (批次大小, 多头注意力的数量, key维度, 序列长度)
        key_t = torch.transpose(key, 2, 3)
        # 张量 query 与 key_t 的矩阵乘法操作对应注意力算法中 query 与 key 做向量点积操作，
        # 其结果 attention_score 的 shape 属性是 (批次大小, 多头注意力的数量, 序列长度, 序列长度)
        attention_score = torch.matmul(query, key_t) / self.sqrt_of_key_dim

        # 通过掩码遮挡用于统一长度的填充标识符以及在解码阶段遮挡未预测信息，
        # 其中编码器自注意力掩码的 shape 属性是 (批次大小, 1, 1, 序列长度),
        # 其中解码器自注意力掩码的 shape 属性是 (批次大小, 1, 序列长度, 序列长度)
        if attention_mask is not None:
            attention_score = torch.masked_fill(attention_score, mask=attention_mask, 
                                                value=constants.MASKED_FILL_VALUE_DEFAULT)

        attention_score = nn.functional.softmax(attention_score, dim=-1)
        # 通过丢弃法防止过拟合
        attention_score = self.dropout(attention_score)
        # 注意力分数与张量 value 的矩阵乘法操作对应注意力算法中注意力分数与 value 做向量点积操作，
        # 其结果 output 的 shape 属性是 (批次大小, 多头注意力的数量, 序列长度, value维度)
        output = torch.matmul(attention_score, value)

        return output, attention_score


class MultiHeadAttention(nn.Module):
    r"""
    多头注意力通过并行地运行多个独立的注意力机制来获取输入序列的不同子空间的注意力分布，
    从而更全面地捕获序列中潜在的多种语义关联。
    """
    
    def __init__(self, head_num: int, embedding_size: int, dropout_prob=constants.MODEL_DROPOUT_PROB_DEFAULT):
        """
        Args:
            head_num: 多头注意力的数量
            embedding_size: 词向量的维度
            dropout_prob: 丢弃率, 防止过拟合
        """
        
        super(MultiHeadAttention, self).__init__()
        self.head_num = head_num
        self.embedding_size = embedding_size
        
        self.query_linear = nn.Linear(embedding_size, embedding_size, bias=False)
        self.key_linear = nn.Linear(embedding_size, embedding_size, bias=False)
        self.value_linear = nn.Linear(embedding_size, embedding_size, bias=False)
        self.out_linear = nn.Linear(embedding_size, embedding_size, bias=False)
        self.layer_norm = nn.LayerNorm(embedding_size, eps=constants.LAYER_NORM_EPS_DEFAULT)

        self.query_dim = embedding_size // head_num
        self.key_dim = embedding_size // head_num
        self.value_dim = embedding_size // head_num
        # 缩放点积注意力
        self.scaled_dot_product_attention = ScaledDotProductAttention(key_dim=self.key_dim, 
                                                                      dropout_prob=dropout_prob)
        
    def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, 
                attention_mask: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        residual = query
        # 批次大小
        batch_size = query.size(0)
        # 序列长度
        sequence_length = query.size(1)
        
        # 操作后 query 的 shape 属性是 (批次大小, 序列长度, 词向量的维度大小)
        query = self.query_linear(query)
        key = self.key_linear(key)
        value = self.value_linear(value)
        
        # 维度修改后 query 的 shape 属性是 (批次大小, 序列长度, 多头注意力的数量, query维度)
        query = query.view(batch_size, -1, self.head_num, self.query_dim)
        # 转置后 query 的 shape 属性是 (批次大小, 多头注意力的数量, 序列长度, query维度)
        query = torch.transpose(query, 1, 2)
        key = key.view(batch_size, -1, self.head_num, self.key_dim)
        key = torch.transpose(key, 1, 2)
        value = value.view(batch_size, -1, self.head_num, self.value_dim)
        value = torch.transpose(value, 1, 2)
        
        if attention_mask is not None:
            # 其中编码器自注意力掩码解压前 attention_mask 的 shape 属性是 (批次大小, 1, 序列长度)，
            # 解压后 attention_mask 的 shape 属性是 (批次大小, 1, 1, 序列长度)。
            # 其中解码器自注意力掩码解压前 attention_mask 的 shape 属性是 (批次大小, 序列长度, 序列长度)，
            # 解压后 attention_mask 的 shape 属性是 (批次大小, 1, 序列长度, 序列长度)。
            attention_mask = torch.unsqueeze(attention_mask, 1)
            
        # 缩放点积注意力返回结果 output 的 shape 属性是 (批次大小, 多头注意力的数量, 序列长度, value维度)
        output, attention_score = self.scaled_dot_product_attention(query, key, value, attention_mask)
        # 转置后 output 的 shape 属性是 (批次大小, 序列长度, 多头注意力的数量, value维度)
        output = torch.transpose(output, 1, 2)
        output = output.contiguous()
        # 维度修改后 output 的 shape 属性是 (批次大小, 序列长度, 词向量的维度大小)
        output = output.view(batch_size, sequence_length, -1)
        
        output = self.out_linear(output)
        # 做残差连接操作
        output = output + residual
        # 做归一化操作
        output = self.layer_norm(output)

        return output, attention_score


class EncoderLayer(nn.Module):
    r"""
    通过堆叠多个相同的编码器层组成编码器，其中每层包括两个子层分别是 多头自注意力、前馈神经网络。
    """
    
    def __init__(self, self_attention: MultiHeadAttention, feed_forward: PositionWiseFeedForward):
        """
        Args:
            decoder_self_attention: 自注意力
            feed_forward: 前馈神经网络
        """
        
        super(EncoderLayer, self).__init__()
        self.self_attention = self_attention
        self.feed_forward = feed_forward
        
    def forward(self, encoder_inputs: torch.Tensor, self_attention_mask: torch.Tensor) \
        -> Tuple[torch.Tensor, torch.Tensor]:
        # 自注意力的返回结果 encoder_outputs 的 shape 属性是 (批次大小, 序列长度, 词向量的维度大小)
        encoder_outputs, attention_score = self.self_attention(query=encoder_inputs, 
                                                               key=encoder_inputs, 
                                                               value=encoder_inputs, 
                                                               attention_mask=self_attention_mask)
        # 前馈神经网络的返回结果 encoder_outputs 的 shape 属性是 (批次大小, 序列长度, 词向量的维度大小)
        encoder_outputs = self.feed_forward(encoder_outputs)
        
        return encoder_outputs, attention_score


class DecoderLayer(nn.Module):
    r"""
    通过堆叠多个相同的解码器层组成编码器，其中每层包括三个子层分别是 多头自注意力、解码器-编码器注意力、前馈神经网络。
    """
    
    def __init__(self, decoder_self_attention: MultiHeadAttention, decoder_encoder_attention: MultiHeadAttention, 
                 feed_forward: PositionWiseFeedForward):
        """
        Args:
            decoder_self_attention: 自注意力
            decoder_encoder_attention: 解码器-编码器注意力
            feed_forward: 前馈神经网络
        """
        
        super(DecoderLayer, self).__init__()
        self.decoder_self_attention = decoder_self_attention
        self.decoder_encoder_attention = decoder_encoder_attention
        self.feed_forward = feed_forward
        
    def forward(self, decoder_inputs: torch.Tensor, encoder_outputs: torch.Tensor, 
                dec_self_attn_mask: torch.Tensor, dec_enc_attn_mask: torch.Tensor) \
                -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        # 自注意力的返回结果 decoder_outputs 的 shape 属性是 (批次大小, 序列长度, 词向量的维度大小)
        decoder_outputs, dec_self_attn_score = self.decoder_self_attention(query=decoder_inputs, 
                                                                           key=decoder_inputs,
                                                                           value=decoder_inputs, 
                                                                           attention_mask=dec_self_attn_mask)
        # 解码器-编码器注意力的返回结果 decoder_outputs 的 shape 属性是 (批次大小, 序列长度, 词向量的维度大小)
        decoder_outputs, dec_enc_attn_score = self.decoder_encoder_attention(query=decoder_outputs, 
                                                                             key=encoder_outputs,
                                                                             value=encoder_outputs, 
                                                                             attention_mask=dec_enc_attn_mask)
        # 前馈神经网络的返回结果 decoder_outputs 的 shape 属性是 (批次大小, 序列长度, 词向量的维度大小)
        decoder_outputs = self.feed_forward(decoder_outputs)
        
        return decoder_outputs, dec_self_attn_score, dec_enc_attn_score


class Encoder(nn.Module):
    r"""
    编码器的任务是将输入序列转换成高维向量表示。
    """
    
    def __init__(self, encoder_layers: nn.ModuleList, word_embedding: nn.Embedding, 
                 positional_encoding: PositionalEncoding, dropout_prob=constants.MODEL_DROPOUT_PROB_DEFAULT):
        """
        Args:
            encoder_layers: 编码器层
            word_embedding: 词向量表
            positional_encoding: 位置编码
            dropout_prob: 丢弃率, 防止过拟合
        """
        
        super(Encoder, self).__init__()
        self.encoder_layers = encoder_layers
        self.word_embedding = word_embedding
        self.positional_encoding = positional_encoding
        self.dropout = nn.Dropout(dropout_prob)
        
    def forward(self, encoder_inputs: torch.Tensor, self_attention_mask: torch.Tensor) \
        -> Tuple[torch.Tensor, torch.Tensor]:
        self_attn_score_list = []
        
        # 词嵌入后变量 encoder_outputs 的 shape 属性是 (批次大小, 序列长度, 词向量的维度大小)
        encoder_outputs = self.word_embedding(encoder_inputs)
        # 添加位置编码
        encoder_outputs = self.positional_encoding(encoder_outputs)
        # 通过丢弃法防止过拟合
        encoder_outputs = self.dropout(encoder_outputs)
        
        # 依次把输出传递到下一个编码器层，其中 encoder_outputs 的 shape 属性是 (批次大小, 序列长度, 词向量的维度大小)
        for encoder_layer in self.encoder_layers:
            encoder_outputs, self_attn_score = encoder_layer(encoder_outputs, self_attention_mask)
            self_attn_score_list.append(self_attn_score)
            
        return encoder_outputs, self_attn_score_list


class Decoder(nn.Module):
    r"""
    解码器的任务是根据编码器生成的高维向量表示和已经生成的目标序列，逐步生成新的序列。
    """
    
    def __init__(self, decoder_layers: nn.ModuleList, word_embedding: nn.Embedding,
                 positional_encoding: PositionalEncoding, dropout_prob=constants.MODEL_DROPOUT_PROB_DEFAULT):
        """
        Args:
            decoder_layers: 编码器层
            word_embedding: 词向量表
            positional_encoding: 位置编码
            dropout_prob: 丢弃率, 防止过拟合
        """
        
        super(Decoder, self).__init__()
        self.decoder_layers = decoder_layers
        self.word_embedding = word_embedding
        self.positional_encoding = positional_encoding
        self.dropout = nn.Dropout(dropout_prob)
        
    def forward(self, decoder_inputs: torch.Tensor, encoder_outputs: torch.Tensor, 
                dec_self_attn_mask: torch.Tensor, dec_enc_attn_mask: torch.Tensor) \
                -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        dec_self_attn_score_list = []
        dec_enc_attn_score_list = []
        
        # 词嵌入后变量 decoder_outputs 的 shape 属性是 (批次大小, 序列长度, 词向量的维度大小)
        decoder_outputs = self.word_embedding(decoder_inputs)
        # 添加位置编码
        decoder_outputs = self.positional_encoding(decoder_outputs)
        # 通过丢弃法防止过拟合
        decoder_outputs = self.dropout(decoder_outputs)
        
        # 依次把输出传递到下一个解码器层，其中 decoder_outputs 的 shape 属性是 (批次大小, 序列长度, 词向量的维度大小)
        for decoder_layer in self.decoder_layers:
            decoder_outputs, dec_self_attn_score, dec_enc_attn_score = decoder_layer(decoder_inputs=decoder_outputs, 
                                                                                     encoder_outputs=encoder_outputs,
                                                                                     dec_self_attn_mask=dec_self_attn_mask,
                                                                                     dec_enc_attn_mask=dec_enc_attn_mask)
            dec_self_attn_score_list.append(dec_self_attn_score)
            dec_enc_attn_score_list.append(dec_enc_attn_score)
            
        return decoder_outputs, dec_self_attn_score_list, dec_enc_attn_score_list


class Transformer(nn.Module):
    r"""
    Transformer模型是基于自注意力机制的深度学习模型。
    """
    
    def __init__(self, encoder: Encoder, decoder: Decoder, output_projection: nn.Module, pad_index: int):
        """
        Args:
            encoder: 编码器
            decoder: 解码器
            output_projection: 输出向量投影到词表空间
            pad_index: 填充标识符索引
        """
        
        super(Transformer, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.output_projection = output_projection
        self.pad_index = pad_index
        
    def forward(self, encoder_inputs: torch.Tensor, decoder_inputs: torch.Tensor, 
                encoder_outputs=None) -> torch.Tensor:
        # 变量 encooder_mask 的 shape 属性是 (批次大小, 1, 序列长度)
        encooder_mask = self.__get_pad_mask(encoder_inputs, self.pad_index)
        # 变量 decoder_mask 的 shape 属性是 (批次大小, 序列长度, 序列长度)
        decoder_mask = self.__get_pad_mask(decoder_inputs, self.pad_index) | self.__get_subsequent_mask(decoder_inputs)
        
        # 用于在推理阶段避免编码器做重复的计算
        if encoder_outputs is None:
            # 编码器输出结果 encoder_outputs 的 shape 属性是 (批次大小, 序列长度, 词向量的维度大小)
            encoder_outputs, *_ = self.encoder(encoder_inputs, encooder_mask)
        
        # 解码器输出结果 encoder_outputs 的 shape 属性是 (批次大小, 序列长度, 词向量的维度大小)
        encoder_outputs, *_ = self.decoder(decoder_inputs, encoder_outputs, decoder_mask, encooder_mask)
        # 投影到词表空间后 decoder_logit 的 shape 属性是 (批次大小, 序列长度, 解码器的词汇表大小)
        decoder_logit = self.output_projection(encoder_outputs)

        return decoder_logit, encoder_outputs

    @staticmethod
    def __get_pad_mask(x: torch.Tensor, pad_index: int) -> torch.Tensor:
        # 变量 mask_tensor 的 shape 属性是 (批次大小, 序列长度)
        mask_tensor = (x == pad_index)
        # 解压后 mask_tensor 的 shape 属性是 (批次大小, 序列长度, 序列长度)
        mask_tensor = mask_tensor.unsqueeze(-2)

        return mask_tensor

    @staticmethod
    def __get_subsequent_mask(x: torch.Tensor) -> torch.Tensor:
        _, sequence_length = x.size()
        ones_tensor = torch.ones(1, sequence_length, sequence_length, device=x.device)
        # 矩阵主对角线及之下区域的元素置为0
        triu_tensor = torch.triu(ones_tensor, diagonal=1)
        triu_tensor = triu_tensor.type(torch.uint8)
        # 变量 mask_tensor 的 shape 属性是 (批次大小, 序列长度, 序列长度)
        mask_tensor = (triu_tensor == 1)
        
        return mask_tensor
